#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
翻译补齐工具
基于Langchain调用OpenAI协议大模型，根据已有翻译补齐缺失字段
"""

import csv
import json
import os
import sys
from typing import List, Dict, Any, Tuple
import yaml
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
from langchain.output_parsers import PydanticOutputParser
from models import TranslationItem, TranslationBatch


class TranslateCompleter:
    """翻译补齐器"""

    def __init__(self, config_path: str = "config.yml"):
        """初始化翻译补齐器"""
        self.config = self._load_config(config_path)
        self.llm = self._init_llm()
        self.parser = PydanticOutputParser(pydantic_object=TranslationBatch)
        self.prompt = self._load_prompt()

    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            # 从环境变量获取API密钥
            if not config['model']['api_key']:
                config['model']['api_key'] = os.getenv('OPENAI_API_KEY')
                if not config['model']['api_key']:
                    raise ValueError("请在config.yml中设置api_key或设置环境变量OPENAI_API_KEY")

            return config
        except FileNotFoundError:
            print(f"错误: 配置文件 {config_path} 不存在")
            sys.exit(1)
        except yaml.YAMLError as e:
            print(f"错误: 配置文件格式错误: {e}")
            sys.exit(1)

    def _init_llm(self) -> ChatOpenAI:
        """初始化语言模型"""
        model_config = self.config['model']
        return ChatOpenAI(
            base_url=model_config['base_url'],
            api_key=model_config['api_key'],
            model=model_config['model_name'],
            temperature=model_config['temperature'],
            max_tokens=model_config['max_tokens']
        )

    def _load_prompt(self) -> str:
        """加载提示文件"""
        prompt_file = self.config['processing']['prompt_file']
        try:
            with open(prompt_file, 'r', encoding='utf-8') as f:
                return f.read().strip()
        except FileNotFoundError:
            print(f"错误: 提示文件 {prompt_file} 不存在")
            sys.exit(1)

    def load_csv_data(self, file_path: str) -> List[Dict[str, str]]:
        """加载CSV数据"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                return list(reader)
        except FileNotFoundError:
            print(f"错误: 输入文件 {file_path} 不存在")
            sys.exit(1)
        except Exception as e:
            print(f"错误: 读取CSV文件失败: {e}")
            sys.exit(1)

    def create_batch_request(self, data_batch: List[Dict[str, str]]) -> str:
        """创建批量请求数据"""
        items = []
        for row in data_batch:
            item = {
                "key": row["key"],
                "cn": row["cn"] if row["cn"] else None,
                "en": row["en"] if row["en"] else None,
                "tr": row["tr"] if row["tr"] else None
            }
            items.append(item)

        return json.dumps(items, ensure_ascii=False, indent=2)

    def process_batch(self, data_batch: List[Dict[str, str]]) -> Tuple[List[TranslationItem], List[Dict[str, str]]]:
        """处理一批数据"""
        # 创建请求数据
        request_data = self.create_batch_request(data_batch)

        # 构建消息
        system_msg = SystemMessage(content=self.prompt)
        human_msg = HumanMessage(content=f"""
请根据以下数据进行翻译补齐：

{request_data}

{self.parser.get_format_instructions()}

请严格按照指定的JSON格式输出结果。
""")

        try:
            # 调用模型
            response = self.llm.invoke([system_msg, human_msg])

            # 解析结果
            result = self.parser.parse(response.content)

            # 检查行数是否匹配
            if len(result.translations) != len(data_batch):
                print(f"警告: 模型返回的行数({len(result.translations)})与输入行数({len(data_batch)})不匹配")
                return [], data_batch

            # 分离有效和无效结果
            valid_items = result.get_valid_items()
            invalid_items = result.get_invalid_items()

            # 对于无效项，返回原始数据
            error_rows = []
            for i, item in enumerate(result.translations):
                if item in invalid_items:
                    error_rows.append(data_batch[i])

            if self.config['logging']['verbose']:
                print(f"批次处理完成: 成功 {len(valid_items)} 项, 失败 {len(error_rows)} 项")

            return valid_items, error_rows

        except Exception as e:
            print(f"处理批次时发生错误: {e}")
            return [], data_batch

    def save_results(self, success_items: List[TranslationItem], error_items: List[Dict[str, str]],
                     append_mode: bool = False):
        """保存结果到文件"""
        # 保存成功的结果
        if success_items:
            output_file = self.config['processing']['output_file']
            mode = 'a' if append_mode else 'w'

            with open(output_file, mode, newline='', encoding='utf-8') as f:
                writer = csv.writer(f, quoting=csv.QUOTE_ALL)

                # 如果是新文件或不是追加模式，写入表头
                if not append_mode or not os.path.exists(output_file) or os.path.getsize(output_file) == 0:
                    writer.writerow(['key', 'cn', 'en', 'tr'])

                for item in success_items:
                    writer.writerow([
                        item.key,
                        item.cn or '',
                        item.en or '',
                        item.tr or ''
                    ])

                # 立即flush到磁盘
                f.flush()
                os.fsync(f.fileno())

            action = "追加" if append_mode else "保存"
            print(f"成功翻译已{action}到 {output_file}, 本次 {len(success_items)} 项")

        # 保存错误的结果
        if error_items:
            error_file = self.config['processing']['error_file']
            error_mode = 'a' if append_mode else 'w'

            with open(error_file, error_mode, newline='', encoding='utf-8') as f:
                writer = csv.writer(f, quoting=csv.QUOTE_ALL)

                # 如果是新文件或不是追加模式，写入表头
                if not append_mode or not os.path.exists(error_file) or os.path.getsize(error_file) == 0:
                    writer.writerow(['key', 'cn', 'en', 'tr'])

                for row in error_items:
                    writer.writerow([row[key] for key in ['key', 'cn', 'en', 'tr']])

                # 立即flush到磁盘
                f.flush()
                os.fsync(f.fileno())

            action = "追加" if append_mode else "保存"
            print(f"错误项目已{action}到 {error_file}, 本次 {len(error_items)} 项")

    def run(self):
        """运行翻译补齐程序"""
        print("开始翻译补齐...")

        # 先清空输出文件
        for file_path in [self.config['processing']['output_file'],
                          self.config['processing']['error_file']]:
            if os.path.exists(file_path):
                os.remove(file_path)

        # 加载数据
        input_file = self.config['processing']['input_file']
        data = self.load_csv_data(input_file)
        print(f"加载了 {len(data)} 行数据")

        # 分批处理
        batch_size = self.config['processing']['batch_size']
        total_success = 0
        total_error = 0

        for i in range(0, len(data), batch_size):
            batch = data[i:i + batch_size]
            batch_num = i//batch_size + 1
            print(f"正在处理第 {batch_num} 批 (行 {i+1}-{min(i+batch_size, len(data))})")

            success_items, error_items = self.process_batch(batch)

            # 每批次结束后立即写入文件
            if success_items or error_items:
                self.save_results(success_items, error_items, append_mode=True)
                total_success += len(success_items)
                total_error += len(error_items)

            print(f"第 {batch_num} 批完成: 成功 {len(success_items)} 项, 失败 {len(error_items)} 项\n")

        print(f"总计: 处理 {len(data)} 行, 成功 {total_success} 行, 失败 {total_error} 行")


def main():
    """主函数"""
    completer = TranslateCompleter()
    completer.run()


if __name__ == "__main__":
    main()